"""
http://baostock.com/baostock/index.php/证券基本资料

create table st_cn_stock_basic_d (
    dt date not null comment '日期',
    stock_code varchar(20) not null default '' comment '股票代码',
    stock_name varchar(30) not null default '' comment '股票中文名',
    ipo_date date comment '上市日期',
    delisting_date date comment '退市日期',
    security_type varchar(10) comment '证券类型 {1:股票, 2:指数, 3:其它, 4:可转债, 5:ETF}',
    listing_status varchar(10) comment '上市状态 {1:上市, 0:退市}',
    PRIMARY KEY (dt, stock_code),
    KEY index_dt (dt)
) ENGINE=InnoDB COMMENT='A股证券基本资料表'

select dt, count(1) from st_cn_stock_basic_d group by dt order by dt;
"""

from common import Args, StockApi, Database
import sys
import asyncio


args = Args.parse()
api = StockApi()
db = Database()

dt = args.date
fields_map = {
    'code': 'stock_code',
    'code_name': 'stock_name',
    'ipoDate': 'ipo_date',
    'outDate': 'delisting_date',
    'type': 'security_type',
    'status': 'listing_status'
}

exist_sql = '''
select count(1) cnt from st_cn_stock_basic_d
where dt=%s
and stock_code=%s
'''


def query_stock_basic(code):
    data = []
    for row in api.query_stock_basic(code=code, fields_map=fields_map): # 这里只报错有基本信息的对象
        row['dt'] = dt
        data.append(row)
        print('processing:\n', row)
        db.execute('delete from st_cn_stock_basic_d where dt=%s and stock_code=%s', (dt, row['stock_code']))
        db.write('st_cn_stock_basic_d', data)


def stock_exists(code):
    cnt = db.execute(exist_sql, (dt, code))[0]['cnt']
    return cnt > 0


stocks = []
for stock in api.query_all_stock(day=dt):
    code = stock['code']
    if code[:2] in ('bj'):
        continue # 跳过不必要的部分
    elif stock_exists(code):
        continue # 跳过已经存在的
    else:
        stocks.append(stock)


async def fn(code, sem):
    async with sem:
        text = await asyncio.get_running_loop().run_in_executor(None, query_stock_basic, code)
 
loop = asyncio.get_event_loop()
sem = asyncio.Semaphore(20)
tasks = [asyncio.ensure_future(fn(i['code'],sem)) for i in stocks]
loop.run_until_complete(asyncio.wait(tasks))
